"""
Author: Morphlng
Date: 2024-03-18 14:08:39
LastEditTime: 2024-12-18 20:31:40
LastEditors: Morphlng
Description: Test for carla adapter api
FilePath: /DrivingGym/tests/carla/carla_adapter_test.py
"""

from driving_gym.simulation.adapter.carla import (
    CarlaActorConfig,
    CarlaConfig,
    CarlaSensorConfig,
)
from driving_gym.simulation.adapter_factory import AdapterFactory
from driving_gym.simulation.common.geom import Location, Rotation, Transform

config = (
    CarlaConfig()
    .connection(host="localhost", port=2000)
    .synchronization(sync_mode=True, fixed_delta_seconds=0.05)
)

adapter = AdapterFactory.create(config.to_dict())


def test_adapter_type():
    assert "CarlaAdapter" in str(type(adapter)), "Adapter type should be CarlaAdapter"


def test_initialize():
    assert adapter._client is None, "Client should be None before initialization"
    adapter.initialize()
    assert adapter._client is not None, "Client should not be None after initialization"


def test_adapter_properties():
    assert adapter._world is None, "Lazy loading world should be None"
    world = adapter.world
    assert adapter._world is not None, "World should be loaded after accessing it"

    assert adapter._map is None, "Lazy loading map should be None"
    map = adapter.map
    assert adapter._map is not None, "Map should be loaded after accessing it"

    assert adapter._hd_map is None, "Lazy loading hd_map should be None"
    hd_map = adapter.hd_map
    assert adapter._hd_map is not None, "HDMap should be loaded after accessing it"


def test_load_world():
    world = adapter.load_world("Town01")
    adapter.tick()
    assert "Town01" in adapter.map.name, "World Town01 should be loaded"

    settings = world.get_settings()
    assert (
        settings.synchronous_mode == config.sync_mode
    ), "Synchronous mode should be turned on"
    assert (
        settings.fixed_delta_seconds == config.fixed_delta_seconds
    ), "Fixed delta seconds should be set correctly"


def test_create_actor():
    actor_config = (
        CarlaActorConfig()
        .blueprint(type="vehicle", blueprint="vehicle.tesla.model3", rolename="hero")
        .spawning(
            spawn_point=Transform(Location(252.4, 195.2, 2.0), Rotation(0, 180, 0)),
            project_to_road=False,
        )
    )

    assert adapter._actor_pool == {}, "Actor pool should be empty"
    adapter.create_actor("hero", actor_config.to_dict())
    assert len(adapter._actor_pool) > 0, "Actor pool should contain the hero actor"

    spawn_trans = adapter.get_transform("hero")
    config_trans = CarlaActorConfig.parse_location(
        actor_config.actor_spawn_point, adapter.map
    )
    assert (
        spawn_trans.location.distance(config_trans.location) < 1.0
    ), "Actor should be spawned at the correct location"


def test_create_sensor():
    sensor_config = (
        CarlaSensorConfig()
        .blueprint(
            type="rgb",
            blueprint="sensor.camera.rgb",
            attributes={
                "image_size_x": 1024,
                "image_size_y": 768,
                "role_name": "front_camera",
            },
        )
        .spawning(parent="hero", attachment_type="spring_arm")
        .position(relative_x=(-2.0, 0.5), relative_z=(2.0, 0.5), relative_pitch=8.0)
    )

    assert adapter._sensor_pool == {}, "Sensor pool should be empty"
    adapter.create_sensor("front_camera", sensor_config.to_dict())
    assert (
        len(adapter._sensor_pool) > 0
    ), "Sensor pool should contain the front_camera sensor"

    for _ in range(10):
        adapter.tick()

    sensor_data = adapter.get_sensor_data("hero:front_camera")
    assert sensor_data is not None, "Sensor callback should be working"

    conf_x = int(sensor_config.sensor_attributes["image_size_x"])
    conf_y = int(sensor_config.sensor_attributes["image_size_y"])
    assert sensor_data.shape == (
        conf_y,
        conf_x,
        3,
    ), f"Image size should be {conf_x}x{conf_y}x3"


def test_sensor_callback():
    from queue import Queue

    queue = Queue()

    sensor_config = (
        CarlaSensorConfig()
        .blueprint(type="rgb")
        .spawning(parent="hero")
        .callback(lambda img: queue.put(img))
    )

    adapter.create_sensor("custom_sensor", sensor_config.to_dict())

    for _ in range(10):
        adapter.tick()

    assert not queue.empty(), "Custom sensor callback should be working"


def test_apply_control():
    move = {"throttle": 1.0, "steer": 0.0, "brake": 0.0, "hand_brake": False}
    for _ in range(10):
        adapter.apply_control("hero", move)
        adapter.tick()

    speed = adapter.get_velocity("hero").length()
    acceleration = adapter.get_acceleration("hero").length()
    assert speed > 0.1, "Vehicle should be moving"
    assert acceleration > 0.1, "Vehicle should be accelerating"

    stop = {"throttle": 0.0, "steer": 0.0, "brake": 1.0, "hand_brake": False}
    for _ in range(10):
        adapter.apply_control("hero", stop)
        adapter.tick()


def test_cleanup():
    client = adapter._client
    actor_pool = adapter._actor_pool.copy()
    sensor_pool = adapter._sensor_pool.copy()

    adapter.cleanup()
    assert adapter._client is None, "Client should be None after cleanup"
    assert adapter._world is None, "World should be None after cleanup"
    assert adapter._map is None, "Map should be None after cleanup"
    assert adapter._hd_map is None, "HDMap should be None after cleanup"
    assert adapter._actor_pool == {}, "Actor pool should be empty after cleanup"
    assert adapter._sensor_pool == {}, "Sensor pool should be empty after cleanup"

    world = client.get_world()
    settings = world.get_settings()
    assert not settings.synchronous_mode, "Synchronous mode should be turned off"

    for actor in list(actor_pool.values()) + [
        sensor.sensor for sensor in sensor_pool.values()
    ]:
        if actor is None:
            continue
        tmp = world.get_actor(actor.id)
        assert (
            tmp is None or not tmp.is_alive or not tmp.is_active
        ), "Actors should be destroyed after cleanup"
